#include <stdio.h>
#include <errno.h>
#include <stdlib.h>
#include <unistd.h>
#include <limits.h>
#include <string.h>
#include <stdbool.h>
#include <util/util.h>
#include <cxl/libcxl.h>
#include <cxl/builtin.h>
#include <fcntl.h>
#include <dirent.h>
#include <sys/stat.h>
#include <stdint.h>
#include <cpuid.h>
#include <signal.h>
#include <numa.h>

#define SYSFS_NODE "/sys/devices/system/node/"

#define MAX_NUMA_NODES (64)
#define MAX_CHAR_LEN (1024)

#define PATH_MSR "/dev/cpu"

#define NODE_CPU (1 << 0)
#define NODE_MEM (1 << 2)

#define OFFSET_MODEL (4)
#define OFFSET_EXT_MODEL (16)

int get_model_id(void)
{
	unsigned int eax = 0, ebx = 0, ecx = 0, edx = 0;
	unsigned int model, ext_model;
	if (__get_cpuid(1, &eax, &ebx, &ecx, &edx) == 0) {
		/* cpuid not supported */
		return 0;
	}
	/* 
		7:4 -  Model
		19:16 - Extended Model ID
		Model ID = Extended Model ID << 4 | Model
	*/
	model = (eax >> OFFSET_MODEL) & 0xf;
	ext_model = (eax >> OFFSET_EXT_MODEL) & 0xf;

	return ext_model << 4 | model;
}

static uint32_t get_msr_offset(void)
{
	unsigned int eax = 0, ebx = 0, ecx = 0, edx = 0;
	char vendor[13] = {
		0,
	};
	unsigned int model;
	if (__get_cpuid(0, &eax, &ebx, &ecx, &edx) == 0) {
		/* cpuid not supported */
		return 0;
	}
	*(unsigned int *)vendor = ebx;
	*(unsigned int *)(vendor + 4) = edx;
	*(unsigned int *)(vendor + 8) = ecx;

	if (!strncmp(vendor, "GenuineIntel", 12)) {
		model = get_model_id();
		if (model == 0xf || model == 0x1d || model == 0x16 ||
		    model == 0x17) /* core 2 family */
			return 0x1a0;
		else
			return 0x1a4;
	} else if (!strncmp(vendor, "AuthenticAMD", 12)) {
		return 0xc0000108;
	} else {
		return 0;
	}
}

static void print_usage_get_latency_matrix(void)
{
	printf("*** cxl get-latency-matrix cmd usage ***\n\n");
	printf("\tcxl get-latency-matrix [Options...]\n\n");
	printf("\t\t   --size <MB>: size(range) of test buffer in MiBs (default: 20000MiB)\n\n");
	printf("\t\t   --stride <Bytes>: stride length in bytes (default: 64B)\n");
	printf("\t\t                     *stride cannot be larger than the size\n\n");
	printf("\t\t   --random: to measure latencies with random access (default: sequential access)\n\n");
	printf("\t\t   --no-change-prefetcher: not to change hw prefetcher before starting test (default: turn-off hw prefetcher before test)\n\n");
	printf("\t\t   --iteration <n>: iterate n times (default: iterate only 1 time)\n\n");
}

double get_read_lat_by_nodes(int node_cpu, int node_mem, unsigned long size,
			     unsigned long stride, int random_access,
			     int num_iter);

static int is_msr_accessible(void)
{
	int fd;
	uint64_t val = 0;
	uint32_t offset = get_msr_offset();

	if ((fd = open("/dev/cpu/0/msr", O_RDWR)) < 0) {
		goto err;
	}
	if (pread(fd, &val, sizeof(val), offset) != sizeof(val)) {
		close(fd);
		goto err;
	}

	if (pwrite(fd, &val, sizeof(val), offset) != sizeof(val)) {
		close(fd);
		goto err;
	}
	close(fd);
	return 1;

err:
	return 0;
}

static int rw_cpus_prefetcher_states(uint64_t *list_val, int is_read)
{
	struct dirent *de;
	struct stat stat_buf;
	DIR *d;
	char filename[MAX_CHAR_LEN];
	uint32_t offset = get_msr_offset();
	if (!offset)
		goto err;

	d = opendir(PATH_MSR);
	if (d) {
		while ((de = readdir(d)) != NULL) {
			if (strcmp(de->d_name, ".") == 0 ||
			    strcmp(de->d_name, "..") == 0)
				continue;
			sprintf(filename, "%s/%s", PATH_MSR, de->d_name);

			if (stat(filename, &stat_buf) == -1)
				goto err;
			if (S_ISDIR(stat_buf.st_mode)) {
				char filename_msr[MAX_CHAR_LEN + 5];
				uint64_t val = 0;
				int cpuid = strtol(de->d_name, NULL, 0);
				int fd;

				sprintf(filename_msr, "%s/msr", filename);
				fd = open(filename_msr, O_RDWR);
				if (fd < 0)
					goto err;

				if (is_read) {
					if (pread(fd, &val, sizeof(val),
						  offset) != sizeof(val)) {
						close(fd);
						goto err;
					}
					list_val[cpuid] = val;
				} else {
					if (list_val[cpuid] == UINT64_MAX) {
						close(fd);
						continue;
					}
					val = list_val[cpuid];

					if (pwrite(fd, &val, sizeof(val),
						   offset) != sizeof(val)) {
						close(fd);
						goto err;
					}
				}
				close(fd);
			}
		}
		closedir(d);
	} else {
		goto err;
	}
	return 0;

err:
	return 1;
}

static uint64_t *orig_val = NULL;
static uint64_t *mod_val = NULL;

static int mod_prefetcher_state(int on)
{
	int num_cpus = sysconf(_SC_NPROCESSORS_CONF);
	orig_val = (uint64_t *)calloc(num_cpus, sizeof(uint64_t));
	if (!orig_val)
		goto err;
	mod_val = (uint64_t *)calloc(num_cpus, sizeof(uint64_t));
	if (!mod_val) {
		free(orig_val);
		goto err;
	}

	for (int i = 0; i < num_cpus; i++)
		orig_val[i] = UINT64_MAX;
	if (rw_cpus_prefetcher_states(orig_val, 1))
		goto err;

	if (on) {
		for (int i = 0; i < num_cpus; i++)
			mod_val[i] = orig_val[i] & (0x0ULL << 0);
	} else {
		for (int i = 0; i < num_cpus; i++)
			mod_val[i] = orig_val[i] | (0xfULL << 0);
	}

	if (rw_cpus_prefetcher_states(mod_val, 0))
		goto err;

	return 0;

err:
	return 1;
}

static void restore_prefetcher_state(void)
{
	if (!orig_val) /* no need to change prefetcher state */
		return;
	if (rw_cpus_prefetcher_states(orig_val, 0))
		error("Cannot change cpus prefetcher state");
	free(orig_val);
	free(mod_val);
	orig_val = mod_val = NULL;

	return;
}

void signal_handler(int sig)
{
	restore_prefetcher_state();
	exit(sig);
}

int get_latency_matrix(int argc, const char **argv)
{
	int num_cpunodes = 0;
	int num_memnodes = 0;
	DIR *d;
	int nodes[MAX_NUMA_NODES] = {
		0,
	};
	int nodes_cpu[MAX_NUMA_NODES] = {
		0,
	};
	int nodes_mem[MAX_NUMA_NODES] = {
		0,
	};
	unsigned long size_mb = 2000;
	unsigned long stride_b = 64;
	int mod_hw_prefetcher = 1;
	int random_access = 0;
	int num_iter = 1;

	/* parse test option */
	for (int i = 1; i < argc; i++) {
		if (!strcmp(argv[i], "--size")) {
			if (i + 1 == argc)
				goto inval_option;
			size_mb = strtoul(argv[++i], NULL, 0);
			if (!size_mb)
				goto inval_option;
		} else if (!strcmp(argv[i], "--stride")) {
			if (i + 1 == argc)
				goto inval_option;
			stride_b = strtoul(argv[++i], NULL, 0);
			if (!stride_b)
				goto inval_option;
		} else if (!strcmp(argv[i], "--no-change-prefetcher")) {
			mod_hw_prefetcher = 0;
		} else if (!strcmp(argv[i], "--random")) {
			random_access = 1;
		} else if (!strcmp(argv[i], "--iteration")) {
			if (i + 1 == argc)
				goto inval_option;
			num_iter = strtoul(argv[++i], NULL, 0);
			if (!num_iter)
				goto inval_option;
		} else {
			goto inval_option;
		}
	}
	if (size_mb * 1024 * 1024 < stride_b)
		goto inval_option;

	/* get online nodes and find init/target */
	d = opendir(SYSFS_NODE);
	if (!d) {
		error("no node sysfs found");
	} else {
		struct dirent *de;
		while ((de = readdir(d)) != NULL) {
			int nd, fd, bytes_read;
			char path_cpulist[MAX_CHAR_LEN];
			char buf[4096];
			if (strncmp(de->d_name, "node", 4))
				continue;
			nd = strtol(de->d_name + 4, NULL, 0);
			sprintf(path_cpulist, "%s/node%d/cpulist", SYSFS_NODE,
				nd);
			fd = open(path_cpulist, O_RDONLY);
			if (fd < 0) {
				error("cannot open %s", path_cpulist);
				continue;
			}
			bytes_read = read(fd, buf, 4096);
			if (bytes_read == 1) //mem node
				nodes[nd] = NODE_MEM;
			else if (bytes_read > 0)
				nodes[nd] = NODE_CPU;
			else
				error("cannot read %s", path_cpulist);
			close(fd);
		}
		closedir(d);
	}

	/* handle SIGINT signal */
	signal(SIGINT, signal_handler);

	/* mod prefetcher state */
	if (mod_hw_prefetcher) {
		if (!is_msr_accessible()) {
			printf("Can't modify prefetcher state. Do 'modprobe msr' with root privileges first,"
			       " then run CXL-CLI with root provileges.\n");
			printf("Enable random_acess to get latency matrix"
			       " without hw prefetcher state modification.\n");
			random_access = 1;
			mod_hw_prefetcher = 0;
		} else {
			mod_prefetcher_state(0);
		}
	}

	/* print perf table header */
	printf("\t    Numa node\t\t(unit: ns)\n");
	printf("Numa node\t");
	for (int i = 0; i < MAX_NUMA_NODES; i++) {
		if (!nodes[i])
			continue;
		if (nodes[i] & NODE_CPU)
			nodes_cpu[num_cpunodes++] = i;
		nodes_mem[num_memnodes++] = i;
		printf("%d\t", i);
	}
	printf("\n");

	/* print actual test results for each node */
	for (int i = 0; i < num_cpunodes; i++) {
		double perf;
		printf("\t%d\t", nodes_cpu[i]);
		fflush(stdout);
		for (int j = 0; j < num_memnodes; j++) {
			perf = get_read_lat_by_nodes(nodes_cpu[i], nodes_mem[j],
						     size_mb * 1024 * 1024,
						     stride_b, random_access,
						     num_iter);
			printf("%.1f\t", perf);
			fflush(stdout);
		}
		printf("\n");
	}

	if (mod_hw_prefetcher)
		restore_prefetcher_state();
	return EXIT_SUCCESS;

inval_option:
	error("Invalid option. Aborting.\n");
	print_usage_get_latency_matrix();
	return EXIT_FAILURE;
}

int cmd_get_latency_matrix(int argc, const char **argv, struct cxl_ctx *ctx)
{
	return get_latency_matrix(argc, argv);
}